from tools.transmit_receive.traffic_template.generator import RandomTemplateValueTemplateGenerator, \
    RandomTimeSpanGenerator
from tools.transmit_receive.traffic_template import TimespanGiver
from tools.transmit_receive import SendUnit
from tools.transmit_receive.transmitter import Transmitter
from tools.o_gan_module import TrafficModelTrainner, TrafficModelUser, TrafficModelParam
import tensorflow as tf
import os

EPOCHS = 10
BATCH_SIZE = 64
NOISE_DIM = 8
TRAFFIC_DIM = 15
G_PATH = 'model/g_model.h5'
D_PATH = 'model/d_model.h5'


def get_or_create_gan_model(g_path: str, d_path: str) -> TrafficModelUser:
    """
    创建一个模型使用者
    搭配上一个训练器
    :param g_path:
    :param d_path:
    :return:
    """
    param = TrafficModelParam(TRAFFIC_DIM, NOISE_DIM)
    # model_trainer = TrafficModelTrainner(EPOCHS, BATCH_SIZE, param)
    return TrafficModelUser(g_path, d_path, param)


def main():
    model_user = get_or_create_gan_model(os.path.join(os.getcwd(), G_PATH),
                                         os.path.join(os.getcwd(), D_PATH))
    val_generator = RandomTemplateValueTemplateGenerator(16)
    val_sequence: list = val_generator.generate()
    # time_generator = RandomTimeSpanGenerator(15)
    # time_sequence: list = time_generator.generate()

    z = model_user.create_noise()
    time_sequence = model_user.generate(z)
    print("z")
    print(z)
    print("time_sequence")
    print(time_sequence)

    timegiver = TimespanGiver()
    sendunits = timegiver.create_units(val_sequence, time_sequence)
    transmitter = Transmitter("192.168.1.110", 6666, sendunits)
    while not transmitter.is_finished():
        sended_data = transmitter.send_data()
        print(sended_data.get_span())
        print(sended_data.get_data())
    print("Over")


if __name__ == '__main__':
    main()
